import numpy as N
import math as ma
import matplotlib.pyplot as plt
from rk_mb import rk4
from Pot import *

class Save():
   def __init__(self):
     pass

memory=Save()
memory.div = 2

def func(y, t): # Dynamical equations, D_y(t) = y(t)
   phi_prime, phi, b_prime, b=y.tolist() # y in
   D_phi_prime = -3.*phi_prime*b_prime/b+dpot_FV(phi) # Equation on phi
   D_phi = phi_prime
   epsilon = memory.epsilon ##
   try : # What was wrong?
     D_b_prime = -epsilon**2*b*(phi_prime**2+pot_FV(phi)) # Equation on b
   except Exception :
     D_b_prime = float('Nan')
   D_b = b_prime 
   result = N.array([ D_phi_prime, D_phi, D_b_prime, D_b ]) # D_y out
   return(result)

def mb_integrate(t_start, t_end, npts, x_start, fun): # Parameters?
   x0=x_start
   #t_list=[0.] # "Time" vector
   #x_list=[N.array([0.,phi0,1.,0.])] # Evolution
   #t_list.append(t_start)
   #x_list.append(x_start)
   t_list=[t_start] # "Time" vector
   x_list=[x_start] # Evolution
   t_vec=N.linspace(t_start,t_end,npts)
   for t0,t1 in zip(t_vec[:-1],t_vec[1:]): # 2 by 2 collection
     dt=t1-t0
     x1=rk4(x0,fun,t0,dt) # Integration by a fourth order Runge-Kutta method
     phi_prime, phi, b_prime, b = x1.tolist()
     if (phi < -3.): # Overshoot
         memory.div = 1
         break # On, if "divide and conquer" on, off otherwise
     elif (phi > 1.): # Undershoot
         memory.div = 0
         break # On, if "divide and conquer" on, off otherwise
     elif (phi_prime > 0.1): # Undershoot followed by oscillations - !!! if memory.div==2, it should be smaller !!!
         memory.div = 0
         break # On, if "divide and conquer" on, off otherwise
     else:
         pass
     t_list .append(t1)
     x_list .append(x1)
     x0=x1
   return((N.array(t_list), N.array(x_list)))

def get_solution_FV(t_start,t_final,phi0,num,epsilon): # Parameters?
   memory.epsilon = epsilon ##
   phi_start = phi0 + (dpot_FV(phi0)/8.)*t_start**2 # phi_start approximation
   dot_phi_start = (dpot_FV(phi0)/4.)*t_start # dot_phi_start approximation
   b_start = t_start # b_start approximation
   dot_b_start = 1. - (epsilon**2*t_start**2/2.)*pot_FV(phi0) # dot_b_start approximation
   y0 = N.array([dot_phi_start, phi_start, dot_b_start, b_start]) # Initial parameters
   result = mb_integrate(t_start, t_final, num, y0, func)
   return(result)

t_start = 0.005 # Start "time" of numerical integration - a small t_start gives a better long-term solution
t_final = 210.5 # "Duration" of the process
pas = 0.005

phi0_M = -7.94071084819e-91 # pas = 0.0005
phi0_m = -7.9407108454e-91

epsilon = 0.01 # On le fixe, apres on cherche phi0
phi0 = (phi0_M+phi0_m)/2.

for i in range(0): # Max number of "divide and conquer"; range(0) for off
    phi0 = (phi0_M+phi0_m)/2.
    tvec,yvec = get_solution_FV(t_start,t_final,phi0,(t_final-t_start)/pas,epsilon)
    if (memory.div == 1):
        phi0_m = phi0
    elif (memory.div == 0):
        phi0_M = phi0
    else:
        break
        
def get_solution_FV_region_II():
    tvec,yvec = get_solution_FV(t_start,t_final,phi0,(t_final-t_start)/pas,epsilon)
    integ = 0.
    f=[integ]
    for i in N.linspace(2,len(tvec),len(tvec)-1):
        vec=tvec[i-2:i].copy()
        bvec=yvec[:,3][i-2:i].copy()
        integ = integ + N.trapz(1./bvec-1./vec,vec) # C'est la methode la plus simple d'integration : utiliser des trapezes
        f.append(integ)
    return((tvec,yvec,f)) # From the vector f[...] I will then construct the function f(...)
